# -*- coding: utf-8 -*-
import contextlib
import geopandas as gpd
from geopandas import GeoDataFrame
from pandas import DataFrame
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from typing import Union, Any
from urllib.parse import quote_plus as urlquote

from config.application import PostgresConfig
from database.model import get_base


class PgManager:
	def __init__(self, pg_config: PostgresConfig) -> None:
		self.pg_config = pg_config
		self.db_connect = f'postgresql+psycopg2://{pg_config.username}:{urlquote(pg_config.password)}@{pg_config.host}:{pg_config.port}/{pg_config.database}'
		self.engine = create_engine(self.db_connect, connect_args={'options': '-csearch_path=public'}, pool_timeout=30)
		self.SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=self.engine)


	@contextlib.contextmanager
	def get_dbsession(self):
		db_session = self.SessionLocal()
		try:
			yield db_session
			db_session.commit()
		except:
			db_session.rollback()
			raise
		finally:
			db_session.close()

	def initial_table(self):
		base = get_base()
		base.metadata.create_all(self.engine)

	def import_shp_to_postgis(self, shp_path: str, table_name: str, srid: int = 4326):
		# 读取shp文件
		gdf: Union[Union[DataFrame, GeoDataFrame], Any] = gpd.read_file(shp_path)
		gdf.crs = f'EPSG:{srid}'
		gdf.to_postgis(table_name, self.engine, if_exists='append')


	#关闭数据库连接
	def close(self):
		self.engine.dispose()




if __name__ == "__main__":
	from database.model import StatisticGraph, StatisticParameter
	from datetime import datetime

	pg_config = PostgresConfig()
	pg_config.host = 'www.rsgisai.com'
	pg_config.port = 15432
	pg_config.database = 'geohazard'
	pg_config.username = 'geohazard'
	pg_config.password = 'cumtb_2023_rerssil'
	pg_manager = PgManager(pg_config)

	with pg_manager.get_dbsession() as db_session:
		a: StatisticParameter = db_session.query(StatisticParameter).filter(
			StatisticParameter.taskId == '56f476d2-e320-4fb2-b462-4cd157232414').first()
		print(a.parameter)
